package com.fanshuai.tree.bst;

import java.util.ArrayList;
import java.util.List;

/**
 * 红黑树
 * 红黑树为二叉搜索树，节点分为红节点和黑节点
 * 1. 根节点和和null节点为黑节点
 * 2. 红节点必须为左节点
 * 3. 不存在左右子结点都为红的结点
 * 4. 不存在父子节点都是红节点
 * 5. 从根节点开始的任何路径，黑节点数量相同
 *
 * 红黑树通过根节点开始的任何路径，黑节点数量相同来保证树的平衡
 * @param <K>
 * @param <V>
 */
public class RedBlackTree<K extends Comparable<K>, V> implements BST<K, V> {
    private boolean RED = true;
    private boolean BLACK = false;

    class Node {
        public K key;
        public V value;
        public boolean color;  //节点颜色
        public Node left;
        public Node right;

        //初始节点为红节点
        public Node(K key, V value) {
            this.key = key;
            this.value = value;
            this.color = RED;
            this.left = null;
            this.right = null;
        }
    }

    private Node root;
    private int size;

    public RedBlackTree() {
        root = null;
        size = 0;
    }

    /**
     * 加入节点
     * 1 按二叉搜索树插入
     * 2 旋转，保持红黑树平衡
     * @param key
     * @param value
     */
    @Override
    public void put(K key, V value) {
        root = put(root, key, value);
    }

    private boolean isRed(Node node) {
        if (null == node) {
            return BLACK;
        }
        return RED;
    }

    /**
     * 左旋  当右孩子为红节点，左孩子不是红节点时，左旋操作保证红节点出现在左子树
     * 右孩子设置为根节点，颜色设置父节点颜色
     * 父节点节点变为左孩子，颜色设置为红色
     * @param node
     * @return
     */
    private Node rotateLeft(Node node) {
        Node x = node.right;
        node.right = x.left;
        x.left = node;

        x.color = node.color;
        node.color = RED;
        return x;
    }

    /**
     * 右旋  当节点的左孩子，和左孩子的左孩子均为红节点时，做右旋操作
     * @param node
     * @return
     */
    private Node rotateRight(Node node) {
        Node x = node.left;
        node.left = x.right;
        x.right = node;

        x.color = node.color;
        node.color = RED;
        return x;
    }

    /**
     * 变色  当左孩子和右孩子均为红节点时，变色为黑节点，父节点转为红节点
     * @param node
     */
    private void flipColors(Node node) {
        node.left.color = BLACK;
        node.right.color = BLACK;
        node.color = RED;
    }

    /**
     * 插入节点
     * 当左孩子和右孩子均为红节点时，做变色操作，及父节点变为红节点，左孩子和右孩子变为黑节点
     * 当新插入叶子节点为右节点时，通过左旋，保证红节点为左节点
     * 当新插入叶子左节点和叶子节点父节点均为红节点时，祖父节点做右旋，保持平衡。右旋后做变色操作
     * @param node
     * @param key
     * @param value
     * @return
     */
    private Node put(Node node, K key, V value) {
        if (null == node) {
            return new Node(key, value);
        }

        if (key.compareTo(node.key) < 0) {
            node.left = put(node.left, key, value);
        } else if (key.compareTo(node.key) > 0) {
            node.right = put(node.right, key, value);
        } else {
            node.value = value;
        }

        if (isRed(node.right) && !isRed(node.left)) { //右孩子为红节点  节点左旋，父节点右旋
            node = rotateLeft(node);
        }
        if (isRed(node.left) && isRed(node.left.left)) { //当节点的左孩子，和左孩子的左孩子均为红节点时，做右旋操作
            node = rotateRight(node);
        }
        if (isRed(node.left) && isRed(node.right)) {//当左孩子和右孩子均为红节点时，做变色操作
            flipColors(node);
        }

        return node;
    }

    @Override
    public V get(K key) {
        if (null == root) {
            return null;
        }

        Node node = root;
        while (node != null) {
            if (node.key.equals(key)) {
                return node.value;
            } else if (key.compareTo(node.key) < 0) {
                node = node.left;
            } else {
                node = node.right;
            }
        }

        return null;
    }

    /**
     * 删除节点
     * @param key
     */
    @Override
    public void delete(K key) {

    }

    @Override
    public List<K> keys() {
        List<K> keys = new ArrayList<>();
        fillKeys(root, keys);
        return keys;
    }

    private void fillKeys(Node node, List<K> keys) {
        if (null == node) {
            return;
        }

        fillKeys(node.left, keys);
        keys.add(node.key);
        fillKeys(node.right, keys);
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean empty() {
        return size == 0;
    }
}
